
import cv2
import torch
import time, math
import numpy as np
import sys

# 计算帧率的变量
fps_time_past = 0
fps_time_now = 0
def init_device():
    device = torch.device("cpu")  # cpu推理模式
    print("-----设备设置成功-----")
    torch_model = torch.hub.load('/Users/a1/PycharmProjects/AimBot/yolov5-6.2', 'custom',
                                 '/Users/a1/PycharmProjects/AimBot/yolov5-6.2/yolov5s.pt', source='local',
                                 force_reload=True)  # 加载本地yolov5模型(需要修改路径和文件)
    print("-----coco模型加载成功-----")
    return torch_model
def camera_init():
    cap = cv2.VideoCapture("/Users/a1/PycharmProjects/AimBot/yolov5-6.2/test3.mp4")#读取视频，如果是使用摄像头的话，需要改为0
    return cap#返回设置
def detect_people(frame,torch_model):
    mode = 0
    people_list = []#重制列表
    fontStyle = cv2.FONT_HERSHEY_TRIPLEX  # 设置显示字体样式
    torch_model = torch_model.to('cpu')#cpu推理模式
    results = torch_model(frame)  # 推理图像
    try:  # 尝试
        xmins = results.pandas().xyxy[0]['xmin']  # 获取所有的xmin坐标
        ymins = results.pandas().xyxy[0]['ymin']  # 获取所有的ymin坐标
        xmaxs = results.pandas().xyxy[0]['xmax']  # 获取所有的xmax坐标
        ymaxs = results.pandas().xyxy[0]['ymax']  # 获取所有的ymax坐标
        class_list = results.pandas().xyxy[0]['class']  # 获取类别信息
        confidences = results.pandas().xyxy[0]['confidence']  # 获取信任度
        count = 0  # 计数
        for xmin, ymin, xmax, ymax, class_l, conf in zip(xmins, ymins, xmaxs, ymaxs, class_list,confidences):  # for循环遍历
            if class_l == 0 and conf >= 0.25:  # 如果置信度大于0.25且识别到的对象为人
                count = count + 1#计数
                people_list.append([int(xmin), int(ymin), int(xmax), int(ymax), conf, count])#将信息加入列表中
            print("识别到的人的数量:",count)
            if len(people_list) > 0:#如果当前帧识别到了人(即列表不为空)
                for i in people_list:#for循环遍历
                    people_x_position = int(i[2] - (i[2] - i[0]) / 2)  # 计算识别目标的中心X坐标
                    people_y_position = int(i[3] - (i[3] - i[1]) / 2)  # 计算识别目标的中心Y坐标
                    if people_y_position>100 and people_y_position <400 and people_x_position>300 and people_x_position<480:#判断识别目标是否处于危险区域内，即学生是否处于危险状态
                        frame = cv2.rectangle(frame, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), (0, 0, 255),2)  # 在原图像中画出方框框记识别目标(红色框)
                        cv2.putText(frame, "Danger", (100, 40), cv2.FONT_HERSHEY_TRIPLEX, 2,
                                    (0, 0, 255))  # 显示帧率在方框之上
                    else: #目标未处于学校危险区域内,即学生处于安全状态
                        frame = cv2.rectangle(frame, (int(i[0]), int(i[1])), (int(i[2]), int(i[3])), (0, 255, 0),2)  # 在原图像中画出方框框记识别目标(绿色框)
        return frame # 返回图像

    except: # 如果在推理过程中出现异常

        return frame # 返回原图像

if __name__  == "__main__" :#主函数运行处

    torch_model = init_device()#初始化推理设备
    cap = camera_init()#初始化视频流或摄像头

    while True:#重复执行
        fps_time_past = time.time()  # 记录开始时间
        ret, frame = cap.read()  # 读取摄像头,获取视频流,ret表示是否读取到视频流,frame表示获取到的图像信息
        frame = detect_people(frame,torch_model)#将获取到的画面送入推理并得到处理后的画面结果
        fps_time_now = time.time()#记录推理完成后的时间
        frame = cv2.rectangle(frame,(300,100),(480,400),(0,0,255),2)#绘制学校内的危险区域
        show_fps_string = "FPS:" + str(int(1 / (fps_time_now - fps_time_past)))  # 计算帧率

        cv2.putText(frame, show_fps_string, (100, 100), cv2.FONT_HERSHEY_TRIPLEX, 2, (0, 255, 255))  # 显示帧率在方框之上
        cv2.imshow('frame',frame)#实时显示推理结果
        if cv2.waitKey(1) & 0xFF == 27:  # 按ESC退出
            print("程序结束")
            break  # 退出
    sys.exit(0)#程序结束